{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac93052e-aa6c-4591-bd81-725585592866",
   "metadata": {},
   "source": [
    "# 混合检索（Vector embdedding + BM25）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d7baef0-d8aa-496e-969a-c12b4439a43e",
   "metadata": {},
   "source": [
    "混合搜索结合了不同搜索范式的优势，以提高检索的准确性和鲁棒性。它既能利用密集向量搜索和稀疏向量搜索的能力，也能利用多种密集向量搜索策略的组合，确保对各种查询进行全面而精确的检索。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b237e39-6cfa-4362-b81a-54276fc9f82b",
   "metadata": {},
   "source": [
    "**混合检索的组成与原理** \n",
    "- 向量检索（语义匹配）    \n",
    "通过将文本转化为高维向量，计算语义相似度召回相关内容。  \n",
    "优势包括：  \n",
    "    - 语义理解：支持复杂语义匹配（如“猫追老鼠”与“猫捕猎老鼠”）  \n",
    "    - 多模态与跨语言：可处理文本、图像、音视频，并支持跨语言检索（如中文查英文内容）  \n",
    "    - 容错性：对拼写错误、模糊描述有较强鲁棒性    \n",
    "局限性：难以精准匹配专有名词（如“iPhone 15”）、缩写（如“RAG”）或ID等短文本。  \n",
    "- 关键词检索（精确匹配）    \n",
    "基于关键词倒排索引或BM25算法，  \n",
    "优势包括：  \n",
    "    - 精准匹配：适合人名、产品名、代码片段等精确查询  \n",
    "    - 低频词捕捉：能识别关键低频词（如“喝咖啡”中的“咖啡”）  \n",
    "局限性：缺乏语义理解能力（例如“北大医院”无法关联“北京大学第一医院”）。  \n",
    "\n",
    "**混合检索的优势**\n",
    "- 覆盖全场景需求  \n",
    "既能处理模糊语义（如用户提问“如何解决失眠”），也能精准匹配特定实体（如“GPT-4的发布时间”）。  \n",
    "- 提升召回率与准确性  \n",
    "实验表明，混合检索相较单一方法可使文档召回率提升20%-30%，尤其适合复杂查询（如含实体与语义混合的提问）。  \n",
    "- 降低噪声干扰  \n",
    "通过重排序技术（如元数据过滤、上下文重排）筛选冗余信息，减少大模型生成时的幻觉风险。  \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed44427-478c-401d-9ecb-806e8cfd5ff8",
   "metadata": {},
   "source": [
    "# Mlivus的混合检索+Rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "349d6676-3fab-40dc-b147-3cffd3f2cf9b",
   "metadata": {},
   "source": [
    "https://milvus.io/docs/zh/milvus_hybrid_search_retriever.md"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e32e2a-4d70-4f23-96ea-89a181dee689",
   "metadata": {},
   "source": [
    "![Milvus](../../../assets/rag9.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d94830fe-f31f-4c37-b414-63a8ab880510",
   "metadata": {},
   "source": [
    "本图展示了最常见的混合搜索方案，即密集+稀疏混合搜索。在这种情况下，使用语义向量相似性和精确关键词匹配两种方法检索候选内容。来自这些方法的结果会被合并、重新排序，并传递给 LLM 以生成最终答案。这种方法兼顾了精确性和语义理解，对各种查询场景都非常有效。  \n",
    "\n",
    "除了密集+稀疏混合搜索，混合策略还可以结合多个密集向量模型。例如，一种密集向量模型可能专门捕捉语义的细微差别，而另一种则侧重于上下文嵌入或特定领域的表示。通过合并这些模型的结果并重新排序，这种类型的混合搜索可确保检索过程更加细致入微、更能感知上下文。  \n",
    "\n",
    "LangChain Milvus集成提供了实现混合搜索的灵活方式，它支持任意数量的向量场，以及任意自定义的密集或稀疏嵌入模型，这使得LangChain Milvus能够灵活适应各种混合搜索使用场景，同时兼容LangChain的其他功能。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "850be57e-9d23-4d0c-9ca8-686ed6467d2b",
   "metadata": {},
   "source": [
    "# 以下内容方便更好理解"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec167c9-819c-4bc1-993c-3c29395abdfd",
   "metadata": {},
   "source": [
    "## 混合检索实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299cfc86-b632-4fc9-bd77-f2cb9478a260",
   "metadata": {},
   "source": [
    "### 固定权重"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7566346b-7d9a-4d77-b3c3-c081839f0d86",
   "metadata": {},
   "source": [
    "获取text用于BM25检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c31bec-ebdd-427c-9da5-aa8d5f42d5dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import connections, Collection\n",
    "\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"192.168.0.188\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "\n",
    "    Args:\n",
    "        collection_name: Milvus 集合名称\n",
    "        host: Milvus 服务器地址（默认 localhost）\n",
    "        port: Milvus 端口（默认 19530）\n",
    "        expr: 过滤条件表达式（默认无过滤）\n",
    "        limit: 返回数据条数上限（默认 1000）\n",
    "        output_fields: 要提取的字段列表（默认 [\"text\"]）\n",
    "\n",
    "    Returns:\n",
    "        list: 包含目标字段值的列表\n",
    "    \"\"\"\n",
    "    # 1. 连接 Milvus\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "\n",
    "    # 2. 加载集合\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "\n",
    "    # 3. 查询数据\n",
    "    results = collection.query(\n",
    "            expr=expr,\n",
    "            output_fields=output_fields,\n",
    "            limit=limit\n",
    "        )\n",
    "\n",
    "    # 4. 提取目标字段为列表\n",
    "    if not output_fields:\n",
    "        raise ValueError(\"output_fields 不能为空\")\n",
    "\n",
    "    field_name = output_fields[0]  # 默认取第一个字段\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "\n",
    "# 示例调用\n",
    "if __name__ == \"__main__\":\n",
    "    # 示例1：读取默认的 text 字段\n",
    "    texts = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    print(f\"获取 {len(texts)} 条文本，前5条: {texts[:5]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71ef7095-536d-4309-9948-61f4290ddfd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama import OllamaLLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "from pymilvus import Collection, connections\n",
    "\n",
    "\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    emb_bgem3 = OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=emb_bgem3,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "def get_llm():\n",
    "    return OllamaLLM(base_url='http://localhost:11434', model='deepseek-r1:14b', temperature=0.1, streaming=True,\n",
    "    callbacks=[StreamingStdOutCallbackHandler()])\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"192.168.0.188\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "    \"\"\"\n",
    "    # 1. 连接 Milvus\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "\n",
    "    # 2. 加载集合\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "\n",
    "    # 3. 查询数据\n",
    "    results = collection.query(\n",
    "            expr=expr,\n",
    "            output_fields=output_fields,\n",
    "            limit=limit\n",
    "        )\n",
    "\n",
    "    # 4. 提取目标字段为列表\n",
    "    if not output_fields:\n",
    "        raise ValueError(\"output_fields 不能为空\")\n",
    "\n",
    "    field_name = output_fields[0]  # 默认取第一个字段\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "def get_qa_chain_with_memory(question: str):\n",
    "    # Initialize memory outside the function so it persists across questions\n",
    "    memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
    "    \n",
    "    vectordb = get_vectordb()\n",
    "\n",
    "    # 1. 初始化 BM25 检索器（关键词检索）\n",
    "    documents = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    bm25_retriever = BM25Retriever.from_texts(documents)\n",
    "    bm25_retriever.k = 10  # 返回前10个BM25检索结果\n",
    "\n",
    "    # 2. 初始化向量检索器\n",
    "    vector_retriever = vectordb.as_retriever(\n",
    "        search_kwargs={\"k\": 10},  # 返回前10个向量检索结果\n",
    "        search_type=\"mmr\",  # 多样性检索\n",
    "    )\n",
    "\n",
    "    # 3. 混合检索（EnsembleRetriever）作为最终检索器\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[0.5, 0.5],  # 调整BM25和向量检索的权重\n",
    "    )\n",
    "\n",
    "    # 4. 定义提示模板\n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 5. 构建对话式检索链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=get_llm(),\n",
    "        retriever=ensemble_retriever,  # 直接使用混合检索器\n",
    "        memory=memory,\n",
    "        output_key=\"answer\",\n",
    "        combine_docs_chain_kwargs={\n",
    "            \"prompt\": QA_PROMPT\n",
    "        },\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result\n",
    "\n",
    "# 测试问题\n",
    "questions = [\n",
    "    \"什么是VMAX？\",\n",
    "    \"VMAX有哪些功能？\",\n",
    "    \"整理成excel表格\"\n",
    "]\n",
    "\n",
    "for question in questions:\n",
    "    result = get_qa_chain_with_memory(question)\n",
    "    print(\"\\n\" + \"=\" * 50 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df084a1b-cf9f-418f-b432-c51675812ead",
   "metadata": {},
   "source": [
    "### 动态权重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ded9e68-6164-41ce-8392-3e9a72f79753",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama import OllamaLLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "from pymilvus import Collection, connections\n",
    "import re\n",
    "\n",
    "# Initialize memory outside the function so it persists across questions\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True\n",
    ")\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    emb_bgem3 = OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=emb_bgem3,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "def get_llm():\n",
    "    return OllamaLLM(base_url='http://localhost:11434', model='deepseek-r1:1.5b', temperature=0.1, streaming=True,\n",
    "    callbacks=[StreamingStdOutCallbackHandler()])\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"192.168.0.188\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "    \"\"\"\n",
    "    # 1. 连接 Milvus\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "\n",
    "    # 2. 加载集合\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "\n",
    "    # 3. 查询数据\n",
    "    results = collection.query(\n",
    "            expr=expr,\n",
    "            output_fields=output_fields,\n",
    "            limit=limit\n",
    "        )\n",
    "\n",
    "    # 4. 提取目标字段为列表\n",
    "    if not output_fields:\n",
    "        raise ValueError(\"output_fields 不能为空\")\n",
    "\n",
    "    field_name = output_fields[0]  # 默认取第一个字段\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "def determine_query_type(question: str) -> str:\n",
    "    \"\"\"\n",
    "    根据问题内容判断查询类型，返回权重调整策略\n",
    "\n",
    "    返回:\n",
    "        \"keyword\" - 更适合关键词检索的问题\n",
    "        \"semantic\" - 更适合语义检索的问题\n",
    "        \"balanced\" - 平衡型问题\n",
    "    \"\"\"\n",
    "    # 关键词型问题特征\n",
    "    keyword_patterns = [\n",
    "        r\"什么是.*\\?\",  # 定义类问题\n",
    "        r\".*包括哪些.*\",  # 列举类问题\n",
    "        r\".*有哪些.*\",  # 列举类问题\n",
    "        r\".*多少种.*\",  # 数量类问题\n",
    "        r\".*步骤.*\",  # 流程类问题\n",
    "        r\".*如何.*\",  # 方法类问题\n",
    "        r\".*怎样.*\",  # 方法类问题\n",
    "        r\".*整理.*表格\",  # 结构化输出要求\n",
    "        r\".*列出.*\",  # 列举要求\n",
    "        r\".*对比.*\",  # 比较类问题\n",
    "    ]\n",
    "\n",
    "    # 语义型问题特征\n",
    "    semantic_patterns = [\n",
    "        r\".*解决.*问题\",  # 解决方案类\n",
    "        r\".*原因.*\",  # 原因分析类\n",
    "        r\".*为什么.*\",  # 原因分析类\n",
    "        r\".*建议.*\",  # 建议类\n",
    "        r\".*优缺点.*\",  # 分析类\n",
    "        r\".*影响.*\",  # 影响分析类\n",
    "        r\".*解释.*\",  # 解释说明类\n",
    "        r\".*理解.*\",  # 理解类\n",
    "        r\".*意味着什么\",  # 含义类\n",
    "    ]\n",
    "\n",
    "    # 检查是否是关键词型问题\n",
    "    for pattern in keyword_patterns:\n",
    "        if re.search(pattern, question):\n",
    "            return \"keyword\"\n",
    "\n",
    "    # 检查是否是语义型问题\n",
    "    for pattern in semantic_patterns:\n",
    "        if re.search(pattern, question):\n",
    "            return \"semantic\"\n",
    "\n",
    "    # 默认平衡型\n",
    "    return \"balanced\"\n",
    "\n",
    "\n",
    "def get_dynamic_weights(query_type: str) -> tuple:\n",
    "    \"\"\"\n",
    "    根据查询类型返回动态权重\n",
    "\n",
    "    返回:\n",
    "        tuple: (bm25_weight, vector_weight)\n",
    "    \"\"\"\n",
    "    if query_type == \"keyword\":\n",
    "        return (0.7, 0.3)  # 更侧重关键词检索\n",
    "    elif query_type == \"semantic\":\n",
    "        return (0.3, 0.7)  # 更侧语义检索\n",
    "    else:\n",
    "        return (0.5, 0.5)  # 平衡权重\n",
    "\n",
    "def get_qa_chain_with_memory(question: str):\n",
    "    vectordb = get_vectordb()\n",
    "\n",
    "    # 1. 确定查询类型和动态权重\n",
    "    query_type = determine_query_type(question)\n",
    "    bm25_weight, vector_weight = get_dynamic_weights(query_type)\n",
    "    print(f\"问题类型: {query_type}, 权重设置: BM25={bm25_weight}, Vector={vector_weight}\")\n",
    "\n",
    "    # 2. 初始化 BM25 检索器\n",
    "    documents = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    bm25_retriever = BM25Retriever.from_texts(documents)\n",
    "    bm25_retriever.k = 10\n",
    "\n",
    "    # 3. 初始化向量检索器\n",
    "    vector_retriever = vectordb.as_retriever(\n",
    "        search_kwargs={\"k\": 10},\n",
    "        search_type=\"mmr\",\n",
    "    )\n",
    "\n",
    "    # 4. 使用动态权重的混合检索\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[bm25_weight, vector_weight],  # 使用动态权重\n",
    "    )\n",
    "    # 5. 混合检索（EnsembleRetriever）作为最终检索器\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[0.5, 0.5],  # 调整BM25和向量检索的权重\n",
    "    )\n",
    "\n",
    "   # 6. 定义提示模板（加入权重信息）\n",
    "   # 4. 定义提示模板\n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 5. 构建对话式检索链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=get_llm(),\n",
    "        retriever=ensemble_retriever,  # 直接使用混合检索器\n",
    "        memory=memory,\n",
    "        output_key=\"answer\",\n",
    "        combine_docs_chain_kwargs={\n",
    "            \"prompt\": QA_PROMPT\n",
    "        },\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result\n",
    "\n",
    "# 测试问题\n",
    "questions = [\n",
    "    \"什么是VMAX的上网日志业务？\",\n",
    "    \"上网日志业务包含哪些功能？\"\n",
    "]\n",
    "\n",
    "for question in questions:\n",
    "    print(f\"\\n问题: {question}\")\n",
    "    result = get_qa_chain_with_memory(question)\n",
    "    print(f\"\\n回答: {result['answer']}\")\n",
    "    print(\"=\" * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5049947-4370-44da-a316-644fd0f8cb6d",
   "metadata": {},
   "source": [
    "## 密集嵌入 + Milvus BM25 内置功能"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7a661e9-d805-4c9b-bcb8-19ba2337365d",
   "metadata": {},
   "source": [
    "https://milvus.io/docs/zh/milvus_hybrid_search_retriever.md"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589d64cd-3dce-471b-9b2e-8d354f1ffab9",
   "metadata": {},
   "source": [
    "## Rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979f1728-da36-4600-8737-d9ccbf251eb1",
   "metadata": {},
   "source": [
    "混合检索结合了向量搜索、关键词匹配等多路径检索方法，能覆盖更广泛的候选文档（高召回率）。但不同检索算法的结果可能存在冗余或噪声，例如：  \n",
    "- 向量搜索因信息压缩（如768维嵌入）导致细节丢失，可能将关键文档排在靠后位置；  \n",
    "- 关键词检索容易受语义鸿沟影响，检索结果可能包含字面匹配但语义无关的内容。  \n",
    "\n",
    "此时，Rerank通过交叉编码器（Cross-Encoder）对混合检索结果进行实时深度语义分析，直接计算查询与文档的匹配度，过滤低相关性内容。例如，原本混合检索返回的10篇文档中，真正相关的可能仅有3篇，Rerank可将其精准筛选并排序，使LLM获得高纯度上下文"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b07ed7f-9dba-4a69-8685-49ffe6ac6550",
   "metadata": {},
   "source": [
    "#### cohere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42bf30d3-7a3d-4b02-b0e5-09ce9682d0e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import CohereRerank\n",
    "# 从环境变量中加载你的 API_KEY\n",
    "_ = load_dotenv(find_dotenv())    # read local .env file\n",
    "cohere_api_key = os.environ['COHERE_API_KEY']\n",
    "\n",
    "# rerank检索\n",
    "# Cohere Rerank配置\n",
    "\n",
    "import cohere\n",
    "cohere_client = cohere.Client(api_key=\"cohere_api_key\")\n",
    "\n",
    "compressor = CohereRerank(\n",
    "    client=cohere_client,\n",
    "    top_n=5,\n",
    "    model=\"rerank-multilingual-v3.0\"  # 支持多语言的新版本\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86feca93-3c19-4936-8add-0a5a2933d553",
   "metadata": {},
   "source": [
    "#### jina"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1764c8-a3ae-4b96-b473-dd5da01df59f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain_community.document_compressors import JinaRerank  # 使用Jina的rerank组件\n",
    "# 从环境变量中加载你的 API_KEY\n",
    "_ = load_dotenv(find_dotenv())    # read local .env file\n",
    "jina_api_key = os.environ['JINA_API_KEY']\n",
    "\n",
    "compressor = JinaRerank(\n",
    "    jina_api_key=jina_api_key,\n",
    "    top_n=3,\n",
    "    model=\"jina-reranker-v2-base-multilingual\"  # Jina的多语言rerank模型[5](@ref)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365fb183-5fd7-459a-9b39-58289452216f",
   "metadata": {},
   "source": [
    "#### BGE版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7330b351-5c7b-4dcb-955c-4ad411198d65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "# from langchain_community.document_compressors import JinaRerank  # 使用Jina的rerank组件\n",
    "\n",
    "# BGE配置\n",
    "# 先将模型下载到本地\n",
    "# modelscope download --model BAAI/bge-reranker-base --cache_dir /opt/workspace/models\n",
    "\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import CrossEncoderReranker\n",
    "from langchain_community.cross_encoders import HuggingFaceCrossEncoder\n",
    "\n",
    "model = HuggingFaceCrossEncoder(model_name=\"/opt/workspace/models/BAAI/bge-reranker-base\")\n",
    "compressor = CrossEncoderReranker(model=model, top_n=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e24c3cd-2804-41d2-93af-17d48c53168b",
   "metadata": {},
   "source": [
    "# 混合检索 + rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b84055fd-600f-4205-950c-47bc0aa295ec",
   "metadata": {},
   "source": [
    "从本地获取BM25检索的DOC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a194f6e-c331-48f1-9a2a-d5df93999e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama import OllamaLLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever\n",
    "from langchain.retrievers.document_compressors import CohereRerank\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "import cohere\n",
    "\n",
    "# Initialize memory outside the function so it persists across questions\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True\n",
    ")\n",
    "\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    emb_bgem3 = OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=emb_bgem3,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "def get_llm():\n",
    "    return OllamaLLM(base_url='http://localhost:11434', model='deepseek-r1:14b', temperature=0.1, streaming=True,\n",
    "    callbacks=[StreamingStdOutCallbackHandler()])\n",
    "\n",
    "def get_qa_chain_with_memory(question: str):\n",
    "    vectordb = get_vectordb()\n",
    "\n",
    "    # 1. 初始化 BM25 检索器（关键词检索）\n",
    "    # 假设你已经有一个文档列表 `documents`（如果没有，可以从 vectordb 获取）\n",
    "    # 示例：documents = vectordb.get_all_documents()\n",
    "    # 这里仅作演示，实际使用时需要替换成你的文档数据\n",
    "    documents = [\"doc1\", \"doc2\", \"doc3\"]  # 替换成你的文档\n",
    "    # 从 Milvus 获取所有原始文本\n",
    "    # documents = get_documents_from_milvus()\n",
    "    bm25_retriever = BM25Retriever.from_texts(documents)\n",
    "    bm25_retriever.k = 10  # 返回前10个BM25检索结果\n",
    "\n",
    "    # 2. 初始化向量检索器\n",
    "    vector_retriever = vectordb.as_retriever(\n",
    "        search_kwargs={\"k\": 10},  # 返回前10个向量检索结果\n",
    "        search_type=\"mmr\",  # 多样性检索\n",
    "    )\n",
    "\n",
    "    # 3. 混合检索（EnsembleRetriever）\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[0.5, 0.5],  # 调整BM25和向量检索的权重\n",
    "    )\n",
    "\n",
    "    # 4. 使用 Jina Rerank 优化结果\n",
    "    # Jina Rerank配置\n",
    "    JINA_API_KEY = \"jina_63bb115e2d5f42d581f42643294792b5CE4nrEINMDcT4vJZJaSLcr5tkbIB\"  # 替换为你的Jina API密钥\n",
    "    \n",
    "    compressor = JinaRerank(\n",
    "        jina_api_key=JINA_API_KEY,\n",
    "        top_n=3,\n",
    "        model=\"jina-reranker-v2-base-multilingual\"  # Jina的多语言rerank模型[5](@ref)\n",
    "    )\n",
    "\n",
    "\n",
    "    compression_retriever = ContextualCompressionRetriever(\n",
    "        base_compressor=compressor,\n",
    "        base_retriever=ensemble_retriever  # 使用混合检索作为基础检索器\n",
    "    )\n",
    "\n",
    "    # 5. 定义提示模板\n",
    "    # 4. 定义提示模板\n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 5. 构建对话式检索链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=get_llm(),\n",
    "        retriever=ensemble_retriever,  # 直接使用混合检索器\n",
    "        memory=memory,\n",
    "        output_key=\"answer\",\n",
    "        combine_docs_chain_kwargs={\n",
    "            \"prompt\": QA_PROMPT\n",
    "        },\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result\n",
    "\n",
    "\n",
    "# 测试问题\n",
    "questions = [\n",
    "    \"什么是VMAX的上网日志业务？\",\n",
    "    \"上网日志业务包含哪些功能？\",\n",
    "    \"整理成excel表格\"\n",
    "]\n",
    "\n",
    "for question in questions:\n",
    "    result = get_qa_chain_with_memory(question)\n",
    "    # print(f\"问题：{question}\")\n",
    "    # print(f\"回答：{result['answer']}\")\n",
    "    # print(\"对话历史：\", memory.load_memory_variables({}))\n",
    "    print(\"\\n\" + \"=\" * 50 + \"\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08518a5b-7ae7-4d26-a238-b33c26b2d1cb",
   "metadata": {},
   "source": [
    "获取milvus数据库中数据text字段"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f221421-3012-4771-9eea-90730e258fdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import connections, Collection\n",
    "\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"129.201.70.35\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "\n",
    "    Args:\n",
    "        collection_name: Milvus 集合名称\n",
    "        host: Milvus 服务器地址（默认 localhost）\n",
    "        port: Milvus 端口（默认 19530）\n",
    "        expr: 过滤条件表达式（默认无过滤）\n",
    "        limit: 返回数据条数上限（默认 1000）\n",
    "        output_fields: 要提取的字段列表（默认 [\"text\"]）\n",
    "\n",
    "    Returns:\n",
    "        list: 包含目标字段值的列表\n",
    "    \"\"\"\n",
    "    # 1. 连接 Milvus\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "\n",
    "    # 2. 加载集合\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "\n",
    "    # 3. 查询数据\n",
    "    results = collection.query(\n",
    "            expr=expr,\n",
    "            output_fields=output_fields,\n",
    "            limit=limit\n",
    "        )\n",
    "\n",
    "    # 4. 提取目标字段为列表\n",
    "    if not output_fields:\n",
    "        raise ValueError(\"output_fields 不能为空\")\n",
    "\n",
    "    field_name = output_fields[0]  # 默认取第一个字段\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "\n",
    "# 示例调用\n",
    "if __name__ == \"__main__\":\n",
    "    # 示例1：读取默认的 text 字段\n",
    "    texts = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    print(f\"获取 {len(texts)} 条文本，前5条: {texts[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c878c584-ee3a-43fa-9299-2a413360c35e",
   "metadata": {},
   "source": [
    "将本地DOC替换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359a34f1-ad7b-47f1-9dbd-2c09862f9131",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama import OllamaLLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever\n",
    "from langchain.retrievers.document_compressors import CohereRerank\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "import cohere\n",
    "\n",
    "# Initialize memory outside the function so it persists across questions\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True\n",
    ")\n",
    "\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    emb_bgem3 = OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=emb_bgem3,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "\n",
    "def get_llm():\n",
    "    return OllamaLLM(base_url='http://localhost:11434', model='deepseek-r1:1.5b', temperature=0.1, streaming=True,\n",
    "    callbacks=[StreamingStdOutCallbackHandler()])\n",
    "\n",
    "\n",
    "from pymilvus import Collection\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"192.168.0.188\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "\n",
    "    Args:\n",
    "        collection_name: Milvus 集合名称\n",
    "        host: Milvus 服务器地址（默认 localhost）\n",
    "        port: Milvus 端口（默认 19530）\n",
    "        expr: 过滤条件表达式（默认无过滤）\n",
    "        limit: 返回数据条数上限（默认 1000）\n",
    "        output_fields: 要提取的字段列表（默认 [\"text\"]）\n",
    "\n",
    "    Returns:\n",
    "        list: 包含目标字段值的列表\n",
    "    \"\"\"\n",
    "    # 1. 连接 Milvus\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "\n",
    "    # 2. 加载集合\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "\n",
    "    # 3. 查询数据\n",
    "    results = collection.query(\n",
    "            expr=expr,\n",
    "            output_fields=output_fields,\n",
    "            limit=limit\n",
    "        )\n",
    "\n",
    "    # 4. 提取目标字段为列表\n",
    "    if not output_fields:\n",
    "        raise ValueError(\"output_fields 不能为空\")\n",
    "\n",
    "    field_name = output_fields[0]  # 默认取第一个字段\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "\n",
    "def get_qa_chain_with_memory(question: str):\n",
    "    vectordb = get_vectordb()\n",
    "\n",
    "    # 1. 初始化 BM25 检索器（关键词检索）\n",
    "    # 假设你已经有一个文档列表 `documents`（如果没有，可以从 vectordb 获取）\n",
    "    # 示例：documents = vectordb.get_all_documents()\n",
    "    # 这里仅作演示，实际使用时需要替换成你的文档数据\n",
    "    # documents = [\"doc1\", \"doc2\", \"doc3\"]  # 替换成你的文档\n",
    "    # 从 Milvus 获取所有原始文本\n",
    "    documents = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    bm25_retriever = BM25Retriever.from_texts(documents)\n",
    "    bm25_retriever.k = 10  # 返回前10个BM25检索结果\n",
    "\n",
    "    # 2. 初始化向量检索器\n",
    "    vector_retriever = vectordb.as_retriever(\n",
    "        search_kwargs={\"k\": 10},  # 返回前10个向量检索结果\n",
    "        search_type=\"mmr\",  # 多样性检索\n",
    "    )\n",
    "\n",
    "    # 3. 混合检索（EnsembleRetriever）\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[0.5, 0.5],  # 调整BM25和向量检索的权重\n",
    "    )\n",
    "\n",
    "    # 4. 使用 Cohere Rerank 优化结果\n",
    "    cohere_client = cohere.Client(api_key=\"Tahx1eySFbKvu9sTyTXrRLf59la3ZUG9vy02stRZ\")\n",
    "    compressor = CohereRerank(\n",
    "        client=cohere_client,\n",
    "        top_n=5,  # 最终保留5个最相关的文档\n",
    "        model=\"rerank-multilingual-v3.0\"\n",
    "    )\n",
    "\n",
    "    compression_retriever = ContextualCompressionRetriever(\n",
    "        base_retriever=ensemble_retriever,  # 使用混合检索作为基础检索器\n",
    "        base_compressor=compressor # 进行rerank\n",
    "\n",
    "    )\n",
    "\n",
    "    # 5. 定义提示模板\n",
    "    # 4. 定义提示模板\n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 5. 构建对话式检索链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=get_llm(),\n",
    "        retriever=ensemble_retriever,  # 直接使用混合检索器\n",
    "        memory=memory,\n",
    "        output_key=\"answer\",\n",
    "        combine_docs_chain_kwargs={\n",
    "            \"prompt\": QA_PROMPT\n",
    "        },\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result\n",
    "\n",
    "\n",
    "# 测试问题\n",
    "questions = [\n",
    "    \"什么是VMAX的上网日志业务？\",\n",
    "    \"上网日志业务包含哪些功能？\",\n",
    "    \"整理成excel表格\"\n",
    "]\n",
    "\n",
    "for question in questions:\n",
    "    result = get_qa_chain_with_memory(question)\n",
    "    # print(f\"问题：{question}\")\n",
    "    # print(f\"回答：{result['answer']}\")\n",
    "    # print(\"对话历史：\", memory.load_memory_variables({}))\n",
    "    print(\"\\n\" + \"=\" * 50 + \"\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4ab52f3-6c48-465b-82ec-f174e2793240",
   "metadata": {},
   "source": [
    "# 混合检索动态权重 + rerank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dda7a5b-dc6a-4e0d-b7af-8547b8a0ecac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama import OllamaLLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever\n",
    "from langchain.retrievers.document_compressors import CohereRerank\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "import cohere\n",
    "from pymilvus import connections, Collection\n",
    "import re\n",
    "\n",
    "# Initialize memory outside the function so it persists across questions\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True\n",
    ")\n",
    "\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    emb_bgem3 = OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=emb_bgem3,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "\n",
    "def get_llm():\n",
    "    return OllamaLLM(base_url='http://localhost:11434', model='deepseek-r1:1.5b', temperature=0.1, streaming=True,\n",
    "                     callbacks=[StreamingStdOutCallbackHandler()])\n",
    "\n",
    "\n",
    "def get_text_list_from_milvus(\n",
    "        collection_name: str,\n",
    "        host: str = \"192.168.0.188\",\n",
    "        port: str = \"19530\",\n",
    "        expr: str = \"\",\n",
    "        limit: int = 1000,\n",
    "        output_fields: list = [\"text\"],\n",
    ") -> list:\n",
    "    \"\"\"\n",
    "    从 Milvus 集合中读取指定字段（默认是 text）并返回列表\n",
    "    \"\"\"\n",
    "    connections.connect(alias=\"default\", host=host, port=port)\n",
    "    collection = Collection(name=collection_name)\n",
    "    collection.load()\n",
    "    results = collection.query(\n",
    "        expr=expr,\n",
    "        output_fields=output_fields,\n",
    "        limit=limit\n",
    "    )\n",
    "    field_name = output_fields[0]\n",
    "    data_list = [item[field_name] for item in results]\n",
    "    return data_list\n",
    "\n",
    "\n",
    "def determine_query_type(question: str) -> str:\n",
    "    \"\"\"\n",
    "    根据问题内容判断查询类型，返回权重调整策略\n",
    "\n",
    "    返回:\n",
    "        \"keyword\" - 更适合关键词检索的问题\n",
    "        \"semantic\" - 更适合语义检索的问题\n",
    "        \"balanced\" - 平衡型问题\n",
    "    \"\"\"\n",
    "    # 关键词型问题特征\n",
    "    keyword_patterns = [\n",
    "        r\"什么是.*\\?\",  # 定义类问题\n",
    "        r\".*包括哪些.*\",  # 列举类问题\n",
    "        r\".*有哪些.*\",  # 列举类问题\n",
    "        r\".*多少种.*\",  # 数量类问题\n",
    "        r\".*步骤.*\",  # 流程类问题\n",
    "        r\".*如何.*\",  # 方法类问题\n",
    "        r\".*怎样.*\",  # 方法类问题\n",
    "        r\".*整理.*表格\",  # 结构化输出要求\n",
    "        r\".*列出.*\",  # 列举要求\n",
    "        r\".*对比.*\",  # 比较类问题\n",
    "    ]\n",
    "\n",
    "    # 语义型问题特征\n",
    "    semantic_patterns = [\n",
    "        r\".*解决.*问题\",  # 解决方案类\n",
    "        r\".*原因.*\",  # 原因分析类\n",
    "        r\".*为什么.*\",  # 原因分析类\n",
    "        r\".*建议.*\",  # 建议类\n",
    "        r\".*优缺点.*\",  # 分析类\n",
    "        r\".*影响.*\",  # 影响分析类\n",
    "        r\".*解释.*\",  # 解释说明类\n",
    "        r\".*理解.*\",  # 理解类\n",
    "        r\".*意味着什么\",  # 含义类\n",
    "    ]\n",
    "\n",
    "    # 检查是否是关键词型问题\n",
    "    for pattern in keyword_patterns:\n",
    "        if re.search(pattern, question):\n",
    "            return \"keyword\"\n",
    "\n",
    "    # 检查是否是语义型问题\n",
    "    for pattern in semantic_patterns:\n",
    "        if re.search(pattern, question):\n",
    "            return \"semantic\"\n",
    "\n",
    "    # 默认平衡型\n",
    "    return \"balanced\"\n",
    "\n",
    "\n",
    "def get_dynamic_weights(query_type: str) -> tuple:\n",
    "    \"\"\"\n",
    "    根据查询类型返回动态权重\n",
    "\n",
    "    返回:\n",
    "        tuple: (bm25_weight, vector_weight)\n",
    "    \"\"\"\n",
    "    if query_type == \"keyword\":\n",
    "        return (0.7, 0.3)  # 更侧重关键词检索\n",
    "    elif query_type == \"semantic\":\n",
    "        return (0.3, 0.7)  # 更侧语义检索\n",
    "    else:\n",
    "        return (0.5, 0.5)  # 平衡权重\n",
    "\n",
    "\n",
    "def get_qa_chain_with_memory(question: str):\n",
    "    vectordb = get_vectordb()\n",
    "\n",
    "    # 1. 确定查询类型和动态权重\n",
    "    query_type = determine_query_type(question)\n",
    "    bm25_weight, vector_weight = get_dynamic_weights(query_type)\n",
    "    print(f\"问题类型: {query_type}, 权重设置: BM25={bm25_weight}, Vector={vector_weight}\")\n",
    "\n",
    "    # 2. 初始化 BM25 检索器\n",
    "    documents = get_text_list_from_milvus(collection_name=\"Vmaxs\")\n",
    "    bm25_retriever = BM25Retriever.from_texts(documents)\n",
    "    bm25_retriever.k = 10\n",
    "\n",
    "    # 3. 初始化向量检索器\n",
    "    vector_retriever = vectordb.as_retriever(\n",
    "        search_kwargs={\"k\": 10},\n",
    "        search_type=\"mmr\",\n",
    "    )\n",
    "\n",
    "    # 4. 使用动态权重的混合检索\n",
    "    ensemble_retriever = EnsembleRetriever(\n",
    "        retrievers=[bm25_retriever, vector_retriever],\n",
    "        weights=[bm25_weight, vector_weight],  # 使用动态权重\n",
    "    )\n",
    "\n",
    "    # 5. 使用 Cohere Rerank 优化结果\n",
    "    cohere_client = cohere.Client(api_key=\"Tahx1eySFbKvu9sTyTXrRLf59la3ZUG9vy02stRZ\")\n",
    "    compressor = CohereRerank(\n",
    "        client=cohere_client,\n",
    "        top_n=5,\n",
    "        model=\"rerank-multilingual-v3.0\"\n",
    "    )\n",
    "\n",
    "    compression_retriever = ContextualCompressionRetriever(\n",
    "        base_compressor=compressor,\n",
    "        base_retriever=ensemble_retriever\n",
    "    )\n",
    "\n",
    "    # 6. 定义提示模板（加入权重信息）\n",
    "    # 4. 定义提示模板\n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 5. 构建对话式检索链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=get_llm(),\n",
    "        retriever=ensemble_retriever,  # 直接使用混合检索器\n",
    "        memory=memory,\n",
    "        output_key=\"answer\",\n",
    "        combine_docs_chain_kwargs={\n",
    "            \"prompt\": QA_PROMPT\n",
    "        },\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result\n",
    "\n",
    "\n",
    "# 测试问题\n",
    "questions = [\n",
    "    \"什么是VMAX的上网日志业务？\",  # 定义类问题，更适合关键词检索\n",
    "    \"上网日志业务包含哪些功能？\",  # 列举类问题，更适合关键词检索\n",
    "    \"整理成excel表格\",  # 结构化输出要求，更适合关键词检索\n",
    "    \"为什么我的VMAX设备会出现日志丢失问题？\",  # 原因分析类，更适合语义检索\n",
    "    \"如何解决VMAX日志存储空间不足的问题？\",  # 解决方案类，更适合语义检索\n",
    "    \"VMAX-S与其他型号的主要区别是什么？\"  # 平衡型问题\n",
    "]\n",
    "\n",
    "for question in questions:\n",
    "    print(f\"\\n问题: {question}\")\n",
    "    result = get_qa_chain_with_memory(question)\n",
    "    print(f\"\\n回答: {result['answer']}\")\n",
    "    print(\"=\" * 50)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
